package com.sunzm.spark.structstream.state

import java.sql.Timestamp

import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
import org.apache.spark.sql.streaming._
import org.slf4j.{Logger, LoggerFactory}

/**
 *
 * 使用mapGroupsWithState完成自定义状态计算
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-07-05 0:38
 */
object StreamingMapGroupsWithStateDemo {
  private val logger: Logger = LoggerFactory.getLogger(this.getClass.getName.stripSuffix("$"))

  def main(args: Array[String]): Unit = {
    val host = "82.156.210.70"
    val port = 9999

    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[*]")
      .config("spark.default.parallelism", 6)
      .config("spark.sql.shuffle.partitions", 6)
      .getOrCreate()

    import spark.implicits._

    //从一个socket端口创建一个DataFrame
    val lines: DataFrame = spark.readStream
      .format("socket")
      .option("host", host)
      .option("port", port)
      .option("includeTimestamp", true)
      .load()

    // Split the lines into words, treat words as sessionId of events
    //s1001,s1002
    val events: Dataset[Event] = lines
      .as[(String, Timestamp)]
      .flatMap {
        case (line, timestamp) => {
          line.split(",").map(word => Event(sessionId = word, timestamp))
        }
      }

    // Sessionize the events. Track number of events, start and end timestamps of session,
    // and report session updates.
    val sessionUpdates = events
      .groupByKey(event => event.sessionId)
      //第一个泛型：状态的类型，第二个类型：输出数据的类型
      .mapGroupsWithState[SessionInfo, SessionUpdate](GroupStateTimeout.ProcessingTimeTimeout) {

        case (sessionId: String, events: Iterator[Event], state: GroupState[SessionInfo]) =>

          // If timed out, then remove session and send final update
          if (state.hasTimedOut) {
            val finalUpdate =
              SessionUpdate(sessionId, state.get.durationMs, state.get.numEvents, expired = true)
            state.remove()
            finalUpdate
          } else {
            // Update start and end timestamps in session
            val timestamps = events.map(_.timestamp.getTime).toSeq
            val updatedSession = if (state.exists) {
              val oldSession = state.get
              SessionInfo(
                oldSession.numEvents + timestamps.size,
                oldSession.startTimestampMs,
                math.max(oldSession.endTimestampMs, timestamps.max))
            } else {
              SessionInfo(timestamps.size, timestamps.min, timestamps.max)
            }
            state.update(updatedSession)

            // Set timeout such that the session will be expired if no data received for 10 seconds
            state.setTimeoutDuration("10 seconds")
            SessionUpdate(sessionId, state.get.durationMs, state.get.numEvents, expired = false)
          }
      }

    // Start running the query that prints the session updates to the console
    val query = sessionUpdates
      .writeStream
      .outputMode(OutputMode.Update())
      .format("console")
      .start()

    query.awaitTermination()

  }

  /** User-defined data type representing the input events */
  private case class Event(sessionId: String, timestamp: Timestamp)

  /**
   * User-defined data type for storing a session information as state in mapGroupsWithState.
   *
   * @param numEvents        total number of events received in the session
   * @param startTimestampMs timestamp of first event received in the session when it started
   * @param endTimestampMs   timestamp of last event received in the session before it expired
   */
  private case class SessionInfo(
                                  numEvents: Int,
                                  startTimestampMs: Long,
                                  endTimestampMs: Long) {

    /** Duration of the session, between the first and last events */
    def durationMs: Long = endTimestampMs - startTimestampMs
  }

  /**
   * User-defined data type representing the update information returned by mapGroupsWithState.
   *
   * @param id         Id of the session
   * @param durationMs Duration the session was active, that is, from first event to its expiry
   * @param numEvents  Number of events received by the session while it was active
   * @param expired    Is the session active or expired
   */
  private case class SessionUpdate(
                                    id: String,
                                    durationMs: Long,
                                    numEvents: Int,
                                    expired: Boolean)
}
